using LinearAlgebra, Kronecker
using Statistics
using Plots, StatsPlots, LaTeXStrings
using DataFrames, DelimitedFiles
import XLSX

################################# Settings

#Algorithm settings
module params_solver

    x_a_H = 0.00001
    x_b_H = 100000
    iter_H_max = 5000
    xd_tol_H = 0.00001

    step_size = 0.125
    max_iterations = 5000
    start_price_multiplier = 10
    xd_tol = 0.00001

end

################################# Downstream tax

#Policy settings
module params_policy
    
    policy = "dt"

    import XLSX
    using Kronecker
    using Statistics
    filename_policy = "aux_params/specs.xlsx"

    #Dimensions
    df = XLSX.readxlsx(filename_policy)["policy_ur"] 
    nI = size(df[:][2:end,1],1)
    nN = 2 
    df = XLSX.readxlsx(filename_policy)["policy_dt"]
    nC = size(df[:][2:end,2],1) 
    nJ = size(df[:][2:end,2:end],2) 
    
    #Emissions
    df = XLSX.readxlsx("aux_params/parameters.xlsx")["emissions_total"]
    tC_per_t = reshape(mean(df[:][2:end,2:3], dims=1),nC,1)

    #Policy
    df = XLSX.readxlsx(filename_policy)["policy_dt"]
    tax_j_c_unscaled = df[:][2:end,2:end].*tC_per_t
    tax_ij_c_unscaled = kronecker(tax_j_c_unscaled, ones(nI,1))
    tax_i_unscaled = zeros(nI,nC)
    subsidy_i_unscaled = zeros(nI)

end

#Import model parameters
include("aux_functions/params_baseline.jl") 
include("aux_functions/params_calibrated.jl") 

#Imperfect competition  
params_policy.conduct = "ic"
include("aux_functions/scale_scc.jl")
df = readdlm(string("results_ic/lf/p_i_c.csv"),',',Float64)
start_p_ic = reshape(df,params.I*params.C,1)
start_s_ij_c_0 = readdlm(string("results_ic/lf/s_ij_c.csv"),',',Float64)
start_p_ij_c = readdlm(string("results_ic/lf/p_ij_c.csv"),',',Float64)
include("aux_functions/solve_equilibrium.jl")

#Perfect competition  
params_policy.conduct = "pc"
include("aux_functions/scale_scc.jl")
df = readdlm(string("results_pc/lf/p_i_c.csv"),',',Float64)
start_p_ic = reshape(df,params.I*params.C,1)
start_s_ij_c_0 = readdlm(string("results_pc/lf/s_ij_c.csv"),',',Float64)
start_p_ij_c = readdlm(string("results_pc/lf/p_ij_c.csv"),',',Float64)
include("aux_functions/solve_equilibrium.jl")


########Extension with mobility

include("aux_functions/params_mobility.jl")
include("aux_functions/params_calibrated.jl") 

#Imperfect competition  
params_policy.conduct = "ic"
include("aux_functions/scale_scc.jl")
df = readdlm(string("results_ic/lf/mobility/p_i_c.csv"),',',Float64)
start_p_ic = reshape(df,params.I*params.C,1)
start_s_ij_c_0 = readdlm(string("results_ic/lf/mobility/s_ij_c.csv"),',',Float64)
start_p_ij_c = readdlm(string("results_ic/lf/mobility/p_ij_c.csv"),',',Float64)
include("aux_functions/solve_equilibrium.jl")



################################# Upstream regulation 

#Policy settings
module params_policy
    
    policy = "ur"

    import XLSX
    using Kronecker
    using Statistics
    filename_policy = "aux_params/specs.xlsx"

    #Dimensions
    df = XLSX.readxlsx(filename_policy)["policy_ur"] 
    nI = size(df[:][2:end,1],1) 
    nN = 2 
    df = XLSX.readxlsx(filename_policy)["policy_dt"]
    nC = size(df[:][2:end,2],1) 
    nJ = size(df[:][2:end,2:end],2)
    
    #AN scale
    df = XLSX.readxlsx("aux_params/parameters.xlsx")["supply"]
    scale_AN = df[:][2:end,10]

    #LUC emissions per ha
    df = XLSX.readxlsx("aux_params/parameters.xlsx")["emissions_luc"]
    tC_per_ha = df[:][2:end,3]

    #Policy parameters
    tax_j_c_unscaled = zeros(nC,nJ)
    tax_ij_c_unscaled = kronecker(tax_j_c_unscaled, ones(nI,1))
    tax_i_unscaled = zeros(nI,nC)
    df = XLSX.readxlsx(filename_policy)["policy_ur"]
    subsidy_i_unscaled = df[:][2:end,4].*tC_per_ha

end

#Import model parameters
include("aux_functions/params_baseline.jl") 
include("aux_functions/params_calibrated.jl") 

#Imperfect competition  
params_policy.conduct = "ic"
include("aux_functions/scale_scc.jl")
df = readdlm(string("results_ic/lf/p_i_c.csv"),',',Float64)
start_p_ic = reshape(df,params.I*params.C,1)
start_s_ij_c_0 = readdlm(string("results_ic/lf/s_ij_c.csv"),',',Float64)
start_p_ij_c = readdlm(string("results_ic/lf/p_ij_c.csv"),',',Float64)
include("aux_functions/solve_equilibrium.jl")

#Perfect competition  
params_policy.conduct = "pc"
include("aux_functions/scale_scc.jl")
df = readdlm(string("results_pc/lf/p_i_c.csv"),',',Float64)
start_p_ic = reshape(df,params.I*params.C,1)
start_s_ij_c_0 = readdlm(string("results_pc/lf/s_ij_c.csv"),',',Float64)
start_p_ij_c = readdlm(string("results_pc/lf/p_ij_c.csv"),',',Float64)
include("aux_functions/solve_equilibrium.jl")



################################# Downstream tax - tariff 

#Policy settings
module params_policy
    
    policy = "tariff"

    import XLSX
    using Kronecker
    using Statistics
    filename_policy = "aux_params/specs.xlsx"

    #Dimensions
    df = XLSX.readxlsx(filename_policy)["policy_ur"] 
    nI = size(df[:][2:end,1],1) 
    nN = 2 
    df = XLSX.readxlsx(filename_policy)["policy_dt"]
    nC = size(df[:][2:end,2],1) 
    nJ = size(df[:][2:end,2:end],2) 
    
    #Emissions
    df = XLSX.readxlsx("aux_params/parameters.xlsx")["emissions_total"]
    tC_per_t = reshape(median(df[:][2:end,2:3], dims=1),nC,1)

    #Policy
    df = XLSX.readxlsx(filename_policy)["policy_tariff"]
    tax_j_c_unscaled = df[:][2:end,2:end].*tC_per_t
    tax_ij_c_unscaled = kronecker(tax_j_c_unscaled, ones(nI,1))
    tax_i_unscaled = zeros(nI,nC)
    subsidy_i_unscaled = zeros(nI)

end

#Import model parameters
include("aux_functions/params_baseline.jl") 
include("aux_functions/params_calibrated.jl") 

#Imperfect competition  
params_policy.conduct = "ic"
include("aux_functions/scale_scc.jl")
df = readdlm(string("results_ic/lf/p_i_c.csv"),',',Float64)
start_p_ic = reshape(df,params.I*params.C,1)
start_s_ij_c_0 = readdlm(string("results_ic/lf/s_ij_c.csv"),',',Float64)
start_p_ij_c = readdlm(string("results_ic/lf/p_ij_c.csv"),',',Float64)
include("aux_functions/solve_equilibrium.jl")



################################# Downstream tax - targeted by region 

#Solver settings
params_solver.step_size = 0.025
params_solver.max_iterations = 7500
params_solver.start_price_multiplier = 20

#Policy settings
module params_policy
    
    policy = "certr"

    import XLSX
    using Kronecker
    using Statistics
    filename_policy = "aux_params/specs.xlsx"

    #Dimensions
    df = XLSX.readxlsx(filename_policy)["policy_ur"] 
    nI = size(df[:][2:end,1],1) 
    nN = 2 
    df = XLSX.readxlsx(filename_policy)["policy_dt"]
    nC = size(df[:][2:end,2],1) 
    nJ = size(df[:][2:end,2:end],2) 
    
    #Emissions
    df = XLSX.readxlsx("aux_params/parameters.xlsx")["emissions_total_region"]
    tC_per_t = df[:][2:end,2:3]
    df = XLSX.readxlsx(filename_policy)["policy_certr"]
    region_in = reshape(df[:][2:end,4],nI,nC)

    #Policy
    tax_j_c_unscaled = zeros(nC,nJ)
    tax_ij_c_unscaled = repeat(reshape(tC_per_t.*region_in, nI*nC,1),1,nJ)
    tax_i_unscaled = zeros(nI,nC)
    subsidy_i_unscaled = zeros(nI)

end

#Import model parameters
include("aux_functions/params_baseline.jl") 
include("aux_functions/params_calibrated.jl") 

#Imperfect competition  
params_policy.conduct = "ic"
include("aux_functions/scale_scc.jl")
df = readdlm(string("results_ic/lf/p_i_c.csv"),',',Float64)
start_p_ic = reshape(df,params.I*params.C,1)
start_s_ij_c_0 = readdlm(string("results_ic/lf/s_ij_c.csv"),',',Float64)
start_p_ij_c = readdlm(string("results_ic/lf/p_ij_c.csv"),',',Float64)
include("aux_functions/solve_equilibrium.jl")



################################# Downstream tax - targeted by county 

#Solver settings
params_solver.step_size = 0.0075
params_solver.max_iterations = 15000
params_solver.start_price_multiplier = 20

#Policy settings
module params_policy
    
    policy = "certc"

    import XLSX
    using Kronecker
    using Statistics
    filename_policy = "aux_params/specs.xlsx"

    #Dimensions
    df = XLSX.readxlsx(filename_policy)["policy_ur"] 
    nI = size(df[:][2:end,1],1)
    nN = 2 
    df = XLSX.readxlsx(filename_policy)["policy_dt"]
    nC = size(df[:][2:end,2],1)
    nJ = size(df[:][2:end,2:end],2)
    
    #Emissions
    df = XLSX.readxlsx("aux_params/parameters.xlsx")["emissions_total"]
    tC_per_t = df[:][2:end,2:3]
    df = XLSX.readxlsx(filename_policy)["policy_certc"]
    county_in = reshape(df[:][2:end,4],nI,nC)

    #Policy
    tax_j_c_unscaled = zeros(nC,nJ)
    tax_ij_c_unscaled = repeat(reshape(tC_per_t.*county_in, nI*nC,1),1,nJ)
    tax_i_unscaled = zeros(nI,nC)
    subsidy_i_unscaled = zeros(nI)

end

#Import model parameters
include("aux_functions/params_baseline.jl") 
include("aux_functions/params_calibrated.jl") 

#Imperfect competition  
params_policy.conduct = "ic"
include("aux_functions/scale_scc.jl")
df = readdlm(string("results_ic/lf/p_i_c.csv"),',',Float64)
start_p_ic = reshape(df,params.I*params.C,1)
start_s_ij_c_0 = readdlm(string("results_ic/lf/s_ij_c.csv"),',',Float64)
start_p_ij_c = readdlm(string("results_ic/lf/p_ij_c.csv"),',',Float64)
include("aux_functions/solve_equilibrium.jl")



################################# First best 

#Solver settings
params_solver.step_size = 0.1
params_solver.max_iterations = 10000
params_solver.start_price_multiplier = 20

#### First best (perfect competition with Pigouvian taxes/subsidies)
module params_policy

    policy = "fb"

    import XLSX
    using Kronecker
    filename_policy = "aux_params/specs.xlsx"

    #Dimensions
    df = XLSX.readxlsx(filename_policy)["policy_ur"] 
    nI = size(df[:][2:end,1],1)
    nN = 2 
    df = XLSX.readxlsx(filename_policy)["policy_dt"]
    nC = size(df[:][2:end,2],1) 
    nJ = size(df[:][2:end,2:end],2) 
    
    #SCC
    df = XLSX.readxlsx("aux_params/scc.xlsx")["baseline_pc"]
    SCC = df[:][2,3]

    #LUC emissions per ha
    df = XLSX.readxlsx("aux_params/parameters.xlsx")["emissions_luc"]
    tC_per_ha = df[:][2:end,3]

    #NLUC emissions per tonne of output
    df = XLSX.readxlsx("aux_params/parameters.xlsx")["emissions_nluc"]
    tC_per_t =  df[:][2:end,2:3]

    #Policy
    tax_j_c = zeros(nC,nJ)
    tax_ij_c = kronecker(tax_j_c, ones(nI,1))
    tax_i = SCC.*tC_per_t
    subsidy_i = SCC.*tC_per_ha

end
params_policy.conduct = "pc"

#Import model parameters
include("aux_functions/params_baseline.jl") 
include("aux_functions/params_calibrated.jl") 

#Starting points
df = readdlm(string("results_pc/lf/p_i_c.csv"),',',Float64)
start_p_ic = reshape(df,params.I*params.C,1)*5
start_s_ij_c_0 = readdlm(string("results_pc/lf/s_ij_c.csv"),',',Float64)
start_p_ij_c = readdlm(string("results_pc/lf/p_ij_c.csv"),',',Float64)
include("aux_functions/solve_equilibrium.jl")



#### Decentralize first best (under imperfect competition with market-power adjusted output taxes)
module params_policy

    policy = "fb"

    import XLSX
    using Kronecker, DelimitedFiles
    filename_policy = "aux_params/specs.xlsx"

    #Dimensions
    df = XLSX.readxlsx(filename_policy)["policy_ur"] 
    nI = size(df[:][2:end,1],1) 
    nN = 2 
    df = XLSX.readxlsx(filename_policy)["policy_dt"]
    nC = size(df[:][2:end,2],1) 
    nJ = size(df[:][2:end,2:end],2) 

    #SCC
    df = XLSX.readxlsx("aux_params/scc.xlsx")["baseline_pc"]
    SCC = df[:][2,3]

    #LUC emissions per ha
    df = XLSX.readxlsx("aux_params/parameters.xlsx")["emissions_luc"]
    tC_per_ha = df[:][2:end,3]

    #Policy
    tax_j_c = zeros(nC,nJ)
    tax_ij_c = kronecker(tax_j_c, ones(nI,1))
    subsidy_i = SCC.*tC_per_ha

    #Markdown adjusted output tax
    results_folder = string("results_pc/fb/")
    tax_fb = readdlm(string(results_folder,"tax_i.csv"),',',Float64)
    p_fb = readdlm(string(results_folder,"p_i_c_pretax.csv"),',',Float64)
    dQdP_fb = readdlm(string(results_folder,"dQdP.csv"),',',Float64)
    dPdQ_fb = dQdP_fb.^(-1)
    df = XLSX.readxlsx("aux_params/parameters.xlsx")["N"]
    N_firms = df[:][2:end,2:end]
    mu_i_c = ( ones(nI,nC) + dPdQ_fb./N_firms ).^(-1)
    tax_i = tax_fb - p_fb.*( ones(nI,nC)-mu_i_c)

end
params_policy.conduct = "ic"

#Starting points
df = readdlm(string("results_pc/fb/p_i_c_pretax.csv"),',',Float64)
start_p_ic = reshape(df,params.I*params.C,1)
start_s_ij_c_0 = readdlm(string("results_pc/fb/s_ij_c.csv"),',',Float64)
start_p_ij_c = readdlm(string("results_pc/fb/p_ij_c.csv"),',',Float64)
include("aux_functions/solve_equilibrium.jl")